
from PIL import Image
import os,json
from torch.utils.data import Dataset
import warnings  
def pil_loader(path):
    return Image.open(path).convert('RGB')
def get_label_map(dataset_name='GlaucomaFundus'):
        if dataset_name=='GlaucomaFundus':
            return {
                 "advanced glaucoma":2,
                "early glaucoma":1,
                "normal" :0
            }
        elif dataset_name=='APTOS':
            return{
        "normal":0,
        "mild Diabetic Retinopathy":1,
        "moderate Diabetic Retinopathy":2,
        "severe Diabetic Retinopathy":3,
        "proliferative Diabetic Retinopathy":4
    }
        else:
            raise NotImplementedError(f"Unknow dataset name: {dataset_name}")
    
class FundusDataset(Dataset):
    def __init__(self, root, split='train', split_name='official',transform=None,dataset_name='GlaucomaFundus'):
        self.data = []
        self.root = root
        self.transform = transform
        self.txt2label=get_label_map(dataset_name)
        with open(os.path.join(root,'annotations.json')) as f:
            self.data_dict=json.load(f)
        
        with open(os.path.join(root,'split',f'{split_name}.json')) as f:
            split_dict=json.load(f)
            if split not in split_dict:
                warnings.warn(f"Invalid split name: {split} in {split_name}.json, returning empty dataset")
                self.split_list = []
                # raise ValueError(f"Invalid split name: {split} in {split_name}.json")
            self.split_list=split_dict[split]
        self.debug=False
    def __getitem__(self, index):
        image_name = self.split_list[index]
        data = self.data_dict[image_name]
        image_path = os.path.join(self.root, data["image_path"])
        image = pil_loader(image_path)
        if self.transform:
            image = self.transform(image)

        # 如果 diagnosis 或 text 缺失，默认标签为 -1
        label = -1  # 默认值，表示未知标签
        if 'diagnosis' in data and data['diagnosis'] is not None:
            if 'text' in data['diagnosis'] and data['diagnosis']['text'] is not None:
                label = self.txt2label.get(data['diagnosis']['text'], -1)  # 使用 get 方法避免 KeyError
        if self.debug:
            return image, {'label': label, 'image_path': image_path,'original_rel_path':data['original_path'],'crop_info':data['crop_info']}
        return image, {'label': label, 'image_path': image_path}

    def __len__(self):
        return len(self.split_list)
